"""Sequential Operations Examples for SSH Slurm MCP Server"""

import asyncio
from fastmcp import Client
import sys
import os

# Add the src directory to the path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from server import mcp, initialize_managers


async def example_sequential_build_workflow():
    """Example: Sequential build and submit workflow like: cd /path/to/project1 && make build && sbatch run.slurm"""
    print("=== Sequential Build Workflow Example ===")
    
    async with Client(mcp) as client:
        connection_id = "researcher@hpc-cluster.example.com:22"
        project_path = "/home/researcher/projects/simulation1"
        
        print("Method 1: Using chained commands (single execution)")
        # Single command with all operations chained
        sequential_result = await client.call_tool("execute_command", {
            "connection_id": connection_id,
            "command": f"cd {project_path} && make build && sbatch run.slurm",
            "working_dir": "~"
        })
        print(f"Sequential command result: {sequential_result.text}")
        
        print("\nMethod 2: Step-by-step with working directory")
        # Step 1: Change to project directory and build
        build_result = await client.call_tool("execute_command", {
            "connection_id": connection_id,
            "command": "make build",
            "working_dir": project_path
        })
        print(f"Build result: {build_result.text}")
        
        # Step 2: Submit Slurm job from the same directory
        submit_result = await client.call_tool("execute_command", {
            "connection_id": connection_id,
            "command": "sbatch run.slurm",
            "working_dir": project_path
        })
        print(f"Submit result: {submit_result.text}")
        
        print("\nMethod 3: Using dedicated Slurm tools")
        # Alternative: Use dedicated submit_slurm_job with pre-built binaries
        # First ensure we're in the right directory and build
        prep_result = await client.call_tool("execute_command", {
            "connection_id": connection_id,
            "command": "cd /home/researcher/projects/simulation1 && make clean && make build",
            "working_dir": "~"
        })
        print(f"Preparation result: {prep_result.text}")
        
        # Read the Slurm script content
        slurm_script_content = """#!/bin/bash
#SBATCH --job-name=simulation1
#SBATCH --partition=compute
#SBATCH --nodes=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=16G
#SBATCH --time=02:00:00

cd /home/researcher/projects/simulation1
echo "Starting simulation in $(pwd)"
./simulation_binary --input=data.txt --output=results.txt
echo "Simulation completed"
"""
        
        submit_job_result = await client.call_tool("submit_slurm_job", {
            "connection_id": connection_id,
            "script_content": slurm_script_content,
            "job_name": "simulation1",
            "partition": "compute",
            "nodes": 1,
            "cpus_per_task": 8,
            "memory": "16G",
            "time_limit": "02:00:00"
        })
        print(f"Slurm job submit result: {submit_job_result.text}")


async def example_complex_project_workflow():
    """Example: Complex multi-step project workflow with dependencies"""
    print("\n=== Complex Project Workflow Example ===")
    
    async with Client(mcp) as client:
        connection_id = "researcher@hpc-cluster.example.com:22"
        
        # Project setup
        projects = [
            {"name": "preprocessing", "path": "/home/researcher/projects/preprocessing"},
            {"name": "simulation", "path": "/home/researcher/projects/simulation"},
            {"name": "analysis", "path": "/home/researcher/projects/analysis"}
        ]
        
        # Step 1: Build all projects in sequence
        print("1. Building all project components...")
        for project in projects:
            build_cmd = f"cd {project['path']} && make clean && make all"
            build_result = await client.call_tool("execute_command", {
                "connection_id": connection_id,
                "command": build_cmd,
                "working_dir": "~"
            })
            print(f"Built {project['name']}: {build_result.text[:100]}...")
        
        # Step 2: Submit dependent Slurm jobs
        print("\n2. Submitting dependent Slurm jobs...")
        
        # Job 1: Preprocessing
        preprocess_script = f"""#!/bin/bash
#SBATCH --job-name=preprocess
#SBATCH --partition=cpu
#SBATCH --nodes=1
#SBATCH --time=30:00

cd {projects[0]['path']}
./preprocess_data --input=/data/raw --output=/data/processed
"""
        
        preprocess_job = await client.call_tool("submit_slurm_job", {
            "connection_id": connection_id,
            "script_content": preprocess_script,
            "job_name": "preprocess",
            "partition": "cpu",
            "nodes": 1,
            "time_limit": "00:30:00"
        })
        print(f"Preprocessing job submitted: {preprocess_job.text}")
        
        # Extract job ID for dependency (in real implementation, you'd parse this)
        preprocess_job_id = "123456"  # This would be parsed from submit result
        
        # Job 2: Simulation (depends on preprocessing)
        simulation_script = f"""#!/bin/bash
#SBATCH --job-name=simulation
#SBATCH --partition=gpu
#SBATCH --nodes=2
#SBATCH --time=02:00:00
#SBATCH --dependency=afterok:{preprocess_job_id}

cd {projects[1]['path']}
mpirun -np 16 ./simulation --input=/data/processed --output=/data/results
"""
        
        simulation_job = await client.call_tool("submit_slurm_job", {
            "connection_id": connection_id,
            "script_content": simulation_script,
            "job_name": "simulation",
            "partition": "gpu",
            "nodes": 2,
            "time_limit": "02:00:00"
        })
        print(f"Simulation job submitted: {simulation_job.text}")


async def example_persistent_session_workflow():
    """Example: Simulating persistent session for interactive development"""
    print("\n=== Persistent Session Workflow Example ===")
    
    async with Client(mcp) as client:
        connection_id = "researcher@hpc-cluster.example.com:22"
        workspace = "/home/researcher/workspace"
        
        # Setup workspace
        setup_commands = [
            f"mkdir -p {workspace}",
            f"cd {workspace}",
            "module load gcc/11.2",
            "module load openmpi/4.1",
            "export OMP_NUM_THREADS=8"
        ]
        
        # Execute setup as a single command chain
        setup_cmd = " && ".join(setup_commands)
        setup_result = await client.call_tool("execute_command", {
            "connection_id": connection_id,
            "command": setup_cmd,
            "working_dir": "~"
        })
        print(f"Workspace setup: {setup_result.text}")
        
        # Development cycle: edit, compile, test, submit
        dev_cycle_commands = [
            # Edit and compile
            f"cd {workspace}",
            "vim simulation.c",  # Interactive editor (would need special handling)
            "gcc -O3 -fopenmp simulation.c -o simulation",
            # Quick test
            "./simulation --test",
            # If test passes, create and submit job
            "echo '#!/bin/bash\n#SBATCH --job-name=sim_test\n#SBATCH --time=10:00\ncd $PWD\n./simulation --full' > test_job.slurm",
            "sbatch test_job.slurm"
        ]
        
        # For non-interactive commands
        non_interactive_cmd = f"cd {workspace} && gcc -O3 -fopenmp simulation.c -o simulation && ./simulation --test"
        test_result = await client.call_tool("execute_command", {
            "connection_id": connection_id,
            "command": non_interactive_cmd,
            "working_dir": "~"
        })
        print(f"Compile and test: {test_result.text}")


async def main():
    """Run all sequential operation examples"""
    print("SSH Slurm MCP Server - Sequential Operations Examples\n")
    
    # Initialize managers
    initialize_managers()
    
    await example_sequential_build_workflow()
    await example_complex_project_workflow() 
    await example_persistent_session_workflow()


if __name__ == "__main__":
    asyncio.run(main())
